#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
简化版 Geppetto WebSocket 服务器
不依赖 NetPyNE 和 NEURON，用于快速测试
"""

import json
import logging
import random
from typing import Dict, Any, Set
from tornado import web, websocket, ioloop

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class SimpleModel:
    """简化的模型（不依赖 NetPyNE）"""
    
    def __init__(self):
        self.network_data = None
        self.simulation_data = None
    
    def create_network(self):
        """创建模拟网络数据"""
        cells = []
        connections = []
        
        # 创建 20 个 PYR 细胞
        for i in range(20):
            cells.append({
                'gid': i,
                'tags': {'pop': 'PYR_pop', 'cellType': 'PYR'},
                'x': random.uniform(0, 100),
                'y': random.uniform(0, 100),
                'z': random.uniform(0, 100),
                'pop': 'PYR_pop',
                'cellType': 'PYR'
            })
        
        # 创建 5 个 PV 细胞
        for i in range(20, 25):
            cells.append({
                'gid': i,
                'tags': {'pop': 'PV_pop', 'cellType': 'PV'},
                'x': random.uniform(0, 100),
                'y': random.uniform(0, 100),
                'z': random.uniform(0, 100),
                'pop': 'PV_pop',
                'cellType': 'PV'
            })
        
        # 创建随机连接
        for i in range(50):
            pre_gid = random.randint(0, 24)
            post_gid = random.randint(0, 24)
            if pre_gid != post_gid:
                connections.append({
                    'preGid': pre_gid,
                    'postGid': post_gid,
                    'weight': random.uniform(0.001, 0.01),
                    'delay': random.uniform(1, 10),
                    'synMech': 'AMPA' if pre_gid < 20 else 'GABA'
                })
        
        self.network_data = {
            'cells': cells,
            'connections': connections,
            'populations': [
                {'name': 'PYR_pop', 'cellType': 'PYR', 'numCells': 20},
                {'name': 'PV_pop', 'cellType': 'PV', 'numCells': 5}
            ]
        }
        
        return self.network_data
    
    def run_simulation(self):
        """运行模拟仿真"""
        spikes = []
        
        # 生成随机 spikes
        for gid in range(25):
            num_spikes = random.randint(5, 15)
            for _ in range(num_spikes):
                t = random.uniform(0, 1000)
                spikes.append({'t': t, 'gid': gid})
        
        # 按时间排序
        spikes.sort(key=lambda x: x['t'])
        
        self.simulation_data = {
            'spikes': spikes,
            'traces': {
                'V_soma': [random.uniform(-70, -50) for _ in range(100)]
            },
            'lfp': [random.uniform(-0.1, 0.1) for _ in range(100)]
        }
        
        return self.simulation_data


class GeppettoWebSocketHandler(websocket.WebSocketHandler):
    """Geppetto WebSocket 处理器"""
    
    clients: Set['GeppettoWebSocketHandler'] = set()
    model: SimpleModel = None
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.client_id = None
    
    def check_origin(self, origin):
        """允许跨域连接"""
        return True
    
    def open(self):
        """WebSocket 连接打开"""
        self.client_id = id(self)
        GeppettoWebSocketHandler.clients.add(self)
        logger.info(f"客户端连接: {self.client_id}, 当前连接数: {len(self.clients)}")
        
        self.send_message({
            'type': 'connection',
            'status': 'connected',
            'message': 'Geppetto WebSocket 服务器已连接',
            'clientId': self.client_id
        })
    
    def on_message(self, message):
        """接收消息"""
        try:
            data = json.loads(message)
            logger.info(f"收到消息: {data.get('type', 'unknown')}")
            
            msg_type = data.get('type')
            
            if msg_type == 'ping':
                self.handle_ping(data)
            elif msg_type == 'init_model':
                self.handle_init_model(data)
            elif msg_type == 'get_network':
                self.handle_get_network(data)
            elif msg_type == 'run_simulation':
                self.handle_run_simulation(data)
            elif msg_type == 'get_simulation_data':
                self.handle_get_simulation_data(data)
            elif msg_type == 'execute_code':
                self.handle_execute_code(data)
            else:
                self.send_error(f"未知的消息类型: {msg_type}")
        
        except json.JSONDecodeError as e:
            logger.error(f"JSON 解析错误: {e}")
            self.send_error(f"JSON 解析错误: {str(e)}")
        except Exception as e:
            logger.error(f"处理消息错误: {e}", exc_info=True)
            self.send_error(f"处理消息错误: {str(e)}")
    
    def on_close(self):
        """WebSocket 连接关闭"""
        GeppettoWebSocketHandler.clients.discard(self)
        logger.info(f"客户端断开: {self.client_id}, 当前连接数: {len(self.clients)}")
    
    def send_message(self, data: Dict[str, Any]):
        """发送消息给客户端"""
        try:
            message = json.dumps(data)
            self.write_message(message)
        except Exception as e:
            logger.error(f"发送消息错误: {e}")
    
    def send_error(self, error_message: str):
        """发送错误消息"""
        self.send_message({
            'type': 'error',
            'error': error_message
        })
    
    def handle_ping(self, data: Dict[str, Any]):
        """处理 ping 消息"""
        self.send_message({
            'type': 'pong',
            'timestamp': data.get('timestamp')
        })
    
    def handle_init_model(self, data: Dict[str, Any]):
        """初始化模型"""
        try:
            logger.info("初始化简化模型...")
            
            if GeppettoWebSocketHandler.model is None:
                GeppettoWebSocketHandler.model = SimpleModel()
            
            self.send_message({
                'type': 'model_initialized',
                'status': 'success',
                'message': '模型初始化成功（简化版）'
            })
            
            logger.info("模型初始化成功")
        
        except Exception as e:
            logger.error(f"初始化模型错误: {e}", exc_info=True)
            self.send_error(f"初始化模型错误: {str(e)}")
    
    def handle_get_network(self, data: Dict[str, Any]):
        """获取网络数据"""
        try:
            logger.info("获取网络数据...")
            
            if GeppettoWebSocketHandler.model is None:
                self.handle_init_model(data)
            
            network_data = GeppettoWebSocketHandler.model.create_network()
            
            self.send_message({
                'type': 'network_data',
                'status': 'success',
                'data': network_data
            })
            
            logger.info(f"网络数据发送成功: {len(network_data['cells'])} 个细胞")
        
        except Exception as e:
            logger.error(f"获取网络数据错误: {e}", exc_info=True)
            self.send_error(f"获取网络数据错误: {str(e)}")
    
    def handle_run_simulation(self, data: Dict[str, Any]):
        """运行仿真"""
        try:
            logger.info("运行仿真...")
            
            if GeppettoWebSocketHandler.model is None:
                self.handle_init_model(data)
            
            self.send_message({
                'type': 'simulation_started',
                'status': 'running',
                'message': '仿真开始运行...'
            })
            
            GeppettoWebSocketHandler.model.run_simulation()
            
            self.send_message({
                'type': 'simulation_completed',
                'status': 'success',
                'message': '仿真运行完成'
            })
            
            logger.info("仿真运行完成")
        
        except Exception as e:
            logger.error(f"运行仿真错误: {e}", exc_info=True)
            self.send_error(f"运行仿真错误: {str(e)}")
    
    def handle_get_simulation_data(self, data: Dict[str, Any]):
        """获取仿真数据"""
        try:
            logger.info("获取仿真数据...")
            
            if GeppettoWebSocketHandler.model is None or GeppettoWebSocketHandler.model.simulation_data is None:
                self.send_error("仿真数据不可用，请先运行仿真")
                return
            
            simulation_data = GeppettoWebSocketHandler.model.simulation_data
            
            self.send_message({
                'type': 'simulation_data',
                'status': 'success',
                'data': simulation_data
            })
            
            logger.info(f"仿真数据发送成功: {len(simulation_data['spikes'])} 个 spikes")
        
        except Exception as e:
            logger.error(f"获取仿真数据错误: {e}", exc_info=True)
            self.send_error(f"获取仿真数据错误: {str(e)}")
    
    def handle_execute_code(self, data: Dict[str, Any]):
        """执行代码"""
        try:
            code = data.get('code', '')
            if not code:
                self.send_error("代码为空")
                return
            
            logger.info(f"执行代码: {code[:100]}...")
            
            exec_globals = {}
            exec_locals = {}
            
            exec(code, exec_globals, exec_locals)
            
            result = exec_locals.get('result', 'Code executed successfully')
            
            self.send_message({
                'type': 'code_executed',
                'status': 'success',
                'result': str(result)
            })
            
            logger.info("代码执行成功")
        
        except Exception as e:
            logger.error(f"执行代码错误: {e}", exc_info=True)
            self.send_error(f"执行代码错误: {str(e)}")


class HealthCheckHandler(web.RequestHandler):
    """健康检查处理器"""
    
    def get(self):
        self.set_header('Content-Type', 'application/json')
        self.write({
            'status': 'ok',
            'service': 'Geppetto WebSocket Server (Simple)',
            'clients': len(GeppettoWebSocketHandler.clients)
        })


def make_app():
    """创建 Tornado 应用"""
    return web.Application([
        (r'/geppetto', GeppettoWebSocketHandler),
        (r'/health', HealthCheckHandler),
    ])


def start_server(port=8888):
    """启动服务器"""
    app = make_app()
    app.listen(port)
    logger.info("=" * 60)
    logger.info("Geppetto WebSocket 服务器（简化版）")
    logger.info("=" * 60)
    logger.info(f"端口: {port}")
    logger.info(f"WebSocket URL: ws://localhost:{port}/geppetto")
    logger.info(f"Health Check URL: http://localhost:{port}/health")
    logger.info("=" * 60)
    logger.info("按 Ctrl+C 停止服务器")
    logger.info("=" * 60)
    
    try:
        ioloop.IOLoop.current().start()
    except KeyboardInterrupt:
        logger.info("\n服务器关闭")


if __name__ == '__main__':
    start_server(port=8888)

